import pandas as pd 
import numpy as np
 
def calculate_kdj(data, n=9, m1=3, m2=3):
    """
    计算KDJ指标
    :param data: DataFrame，需包含列['open', 'close', 'high', 'low']
    :param n: RSV计算周期（默认9天）
    :param m1: K值平滑周期（默认3）
    :param m2: D值平滑周期（默认3）
    :return: 添加了K、D、J列的DataFrame
    """
    # 初始化结果列 
    data['RSV'] = np.nan  
    data['K'] = np.nan  
    data['D'] = np.nan  
    data['J'] = np.nan  
    
    # 从第n-1天开始计算（保证有足够历史数据）
    for i in range(n-1, len(data)):
        # 1. 计算周期内极值
        low_min = data['low'].iloc[i-n+1:i+1].min()
        high_max = data['high'].iloc[i-n+1:i+1].max()
        
        # 2. 计算RSV（处理分母为0的情况）
        if high_max != low_min:
            rsv = (data['close'].iloc[i] - low_min) / (high_max - low_min) * 100
        else:
            rsv = 50  # 涨停/跌停的特殊处理
        data.loc[data.index[i],  'RSV'] = rsv
        
        # 3. 初始化首日K/D值 
        if i == n-1:
            k = rsv if not np.isnan(rsv)  else 50 
            d = 50 
        else:
            # 4. 递推计算K值和D值 
            prev_k = data['K'].iloc[i-1]
            prev_d = data['D'].iloc[i-1]
            k = (2/3) * prev_k + (1/3) * rsv
            d = (2/3) * prev_d + (1/3) * k
        
        # 5. 计算J值
        j = 3 * k - 2 * d 
        
        # 存储结果 
        data.loc[data.index[i],  'K'] = k 
        data.loc[data.index[i],  'D'] = d 
        data.loc[data.index[i],  'J'] = j 
    
    return data 
 
# 示例用法 
if __name__ == "__main__":
    # 构造示例数据（替换为实际数据）
    dates = pd.date_range(start="2025-10-01",  periods=30, freq="D")
    np.random.seed(42) 
    data = pd.DataFrame({
        'open': np.random.uniform(90,  110, 30).round(2),
        'close': np.random.uniform(95,  105, 30).round(2),
        'high': np.random.uniform(100,  120, 30).round(2),
        'low': np.random.uniform(85,  100, 30).round(2)
    }, index=dates)
    
    # 计算KDJ指标
    result = calculate_kdj(data)
    
    # 查看结果（最后5天）
    print(result[['close', 'K', 'D', 'J']].tail())